package com.fanxuankai.kafka.connect.transforms;

import com.google.common.base.Strings;
import org.apache.kafka.common.config.ConfigDef;
import org.apache.kafka.connect.connector.ConnectRecord;
import org.apache.kafka.connect.data.ConnectSchema;
import org.apache.kafka.connect.data.Field;
import org.apache.kafka.connect.data.Schema;
import org.apache.kafka.connect.data.Struct;
import org.apache.kafka.connect.transforms.ExtractField;
import org.apache.kafka.connect.transforms.Transformation;
import org.apache.kafka.connect.transforms.util.SimpleConfig;

import java.util.Map;

import static org.apache.kafka.connect.transforms.util.Requirements.requireMapOrNull;
import static org.apache.kafka.connect.transforms.util.Requirements.requireStructOrNull;

/**
 * 抽取 key 字段时，不指定则默认为原先的 key，其他功能不变
 *
 * @author fanxuankai
 * @see ExtractField
 */
public abstract class ExtractFieldAuto<R extends ConnectRecord<R>> implements Transformation<R> {

    public static final String OVERVIEW_DOC =
            "Extract the specified field from a Struct when schema present, or a " + "Map in the case of schemaless " +
                    "data. " + "Any null values are passed through unmodified." + "<p/>Use the " + "concrete " +
                    "transformation type designed for the record key (<code>" + org.apache.kafka.connect.transforms.ExtractField.Key.class.getName() + "</code>) " + "or value (<code>" + org.apache.kafka.connect.transforms.ExtractField.Value.class.getName() + "</code>).";

    private static final String FIELD_CONFIG = "field";
    private static final String NAME_CONFIG = "name";

    public static final ConfigDef CONFIG_DEF = new ConfigDef()
            .define(FIELD_CONFIG, ConfigDef.Type.STRING, ConfigDef.NO_DEFAULT_VALUE,
                    ConfigDef.Importance.MEDIUM, "Field name to extract.")
            .define(NAME_CONFIG, ConfigDef.Type.STRING,
                    ConfigDef.NO_DEFAULT_VALUE, ConfigDef.Importance.MEDIUM, "Key Field name.");

    private static final String PURPOSE = "field extraction";

    private String sourceFieldName;
    private String targetFieldName;

    @Override
    public void configure(Map<String, ?> props) {
        final SimpleConfig config = new SimpleConfig(CONFIG_DEF, props);
        sourceFieldName = config.getString(FIELD_CONFIG);
        targetFieldName = config.getString(NAME_CONFIG);
    }

    @Override
    public R apply(R record) {
        String fieldName = this.sourceFieldName;
        boolean fieldNameIsEmpty = fieldName == null || fieldName.isEmpty();
        if (fieldNameIsEmpty && getClass() == Key.class) {
            if (record.keySchema() == null
                    || record.keySchema().type() != Schema.Type.STRUCT
                    || record.keySchema().fields().size() > 1) {
                return record;
            }
            fieldName = record.keySchema().fields().get(0).name();
        }
        final Schema schema = operatingSchema(record);
        if (schema == null) {
            final Map<String, Object> value = requireMapOrNull(operatingValue(record), PURPOSE);
            return newRecord(record, null, value == null ? null : value.get(fieldName));
        } else {
            final Struct value = requireStructOrNull(operatingValue(record), PURPOSE);
            Schema targetSchema;
            if (schema.type() == Schema.Type.STRUCT) {
                Field field = schema.field(fieldName);
                if (field == null) {
                    throw new IllegalArgumentException("Unknown field: " + fieldName);
                }
                if (Strings.isNullOrEmpty(targetFieldName)) {
                    targetSchema = field.schema();
                } else {
                    targetSchema = new ConnectSchema(schema.type(), schema.isOptional(), schema.defaultValue(),
                            targetFieldName, schema.version(), schema.doc(), schema.parameters(),
                            schema.fields(), schema.keySchema(), schema.valueSchema());
                }
            } else {
                if (Strings.isNullOrEmpty(targetFieldName)) {
                    targetSchema = schema.keySchema();
                } else {
                    targetSchema = new ConnectSchema(schema.type(), schema.isOptional(), schema.defaultValue(),
                            targetFieldName, schema.version(), schema.doc(), schema.parameters(),
                            null, schema.keySchema(), schema.valueSchema());
                }
            }
            return newRecord(record, targetSchema, value == null ? null : value.get(fieldName));
        }
    }

    @Override
    public void close() {
    }

    @Override
    public ConfigDef config() {
        return CONFIG_DEF;
    }

    protected abstract Schema operatingSchema(R record);

    protected abstract Object operatingValue(R record);

    protected abstract R newRecord(R record, Schema updatedSchema, Object updatedValue);

    public static class Key<R extends ConnectRecord<R>> extends ExtractFieldAuto<R> {
        @Override
        protected Schema operatingSchema(R record) {
            return record.keySchema();
        }

        @Override
        protected Object operatingValue(R record) {
            return record.key();
        }

        @Override
        protected R newRecord(R record, Schema updatedSchema, Object updatedValue) {
            return record.newRecord(record.topic(), record.kafkaPartition(), updatedSchema, updatedValue,
                    record.valueSchema(), record.value(), record.timestamp());
        }
    }

    public static class Value<R extends ConnectRecord<R>> extends ExtractFieldAuto<R> {
        @Override
        protected Schema operatingSchema(R record) {
            return record.valueSchema();
        }

        @Override
        protected Object operatingValue(R record) {
            return record.value();
        }

        @Override
        protected R newRecord(R record, Schema updatedSchema, Object updatedValue) {
            return record.newRecord(record.topic(), record.kafkaPartition(), record.keySchema(), record.key(),
                    updatedSchema, updatedValue, record.timestamp());
        }
    }
}
